{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "\n",
    "\n",
    "class SMPLModel():\n",
    "    def __init__(self, model_path):\n",
    "        \"\"\"\n",
    "        SMPL model.\n",
    "    \n",
    "        Parameter:\n",
    "        ---------\n",
    "        model_path: Path to the SMPL model parameters, pre-processed by\n",
    "        `preprocess.py`.\n",
    "    \n",
    "        \"\"\"\n",
    "        with open(model_path, 'rb') as f:\n",
    "            params = pickle.load(f)\n",
    "\n",
    "            self.J_regressor = params['J_regressor']\n",
    "            self.weights = np.asarray(params['weights'])\n",
    "            self.posedirs = np.asarray(params['posedirs'])\n",
    "            self.v_template = np.asarray(params['v_template'])\n",
    "            self.shapedirs = np.asarray(params['shapedirs'])\n",
    "            self.faces = np.asarray(params['f'])\n",
    "            self.kintree_table = np.asarray(params['kintree_table'])\n",
    "\n",
    "        id_to_col = {\n",
    "            self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])\n",
    "        }\n",
    "        self.parent = {\n",
    "            i: id_to_col[self.kintree_table[0, i]]\n",
    "            for i in range(1, self.kintree_table.shape[1])\n",
    "        }\n",
    "\n",
    "        self.pose_shape = [24, 3]\n",
    "        self.beta_shape = [10]\n",
    "        self.trans_shape = [3]\n",
    "\n",
    "        self.pose = np.zeros(self.pose_shape)\n",
    "        self.beta = np.zeros(self.beta_shape)\n",
    "        self.trans = np.zeros(self.trans_shape)\n",
    "\n",
    "        self.verts = None\n",
    "        self.J = None\n",
    "        self.R = None\n",
    "        self.G = None\n",
    "\n",
    "        self.update()\n",
    "\n",
    "    def set_params(self, pose=None, beta=None, trans=None):\n",
    "        \"\"\"\n",
    "        Set pose, shape, and/or translation parameters of SMPL model. Verices of the\n",
    "        model will be updated and returned.\n",
    "    \n",
    "        Prameters:\n",
    "        ---------\n",
    "        pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation\n",
    "        relative to parent joint. For root joint it's global orientation.\n",
    "        Represented in a axis-angle format.\n",
    "    \n",
    "        beta: Parameter for model shape. A vector of shape [10]. Coefficients for\n",
    "        PCA component. Only 10 components were released by MPI.\n",
    "    \n",
    "        trans: Global translation of shape [3].\n",
    "    \n",
    "        Return:\n",
    "        ------\n",
    "        Updated vertices.\n",
    "    \n",
    "        \"\"\"\n",
    "        if pose is not None:\n",
    "            self.pose = pose\n",
    "        if beta is not None:\n",
    "            self.beta = beta\n",
    "        if trans is not None:\n",
    "            self.trans = trans\n",
    "        self.update()\n",
    "        return self.verts\n",
    "\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Called automatically when parameters are updated.\n",
    "    \n",
    "        \"\"\"\n",
    "        # how beta affect body shape\n",
    "        v_shaped = self.shapedirs.dot(self.beta) + self.v_template\n",
    "        # joints location\n",
    "        self.J = self.J_regressor.dot(v_shaped)\n",
    "        pose_cube = self.pose.reshape((-1, 1, 3))\n",
    "        # rotation matrix for each joint\n",
    "        self.R = self.rodrigues(pose_cube)\n",
    "        I_cube = np.broadcast_to(\n",
    "            np.expand_dims(np.eye(3), axis=0),\n",
    "            (self.R.shape[0] - 1, 3, 3)\n",
    "        )\n",
    "        lrotmin = (self.R[1:] - I_cube).ravel()\n",
    "        # how pose affect body shape in zero pose\n",
    "        v_posed = v_shaped + self.posedirs.dot(lrotmin)\n",
    "        # world transformation of each joint\n",
    "        G = np.empty((self.kintree_table.shape[1], 4, 4))\n",
    "        G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))\n",
    "        for i in range(1, self.kintree_table.shape[1]):\n",
    "            G[i] = G[self.parent[i]].dot(\n",
    "                self.with_zeros(\n",
    "                    np.hstack(\n",
    "                        [self.R[i], ((self.J[i, :] - self.J[self.parent[i], :]).reshape([3, 1]))]\n",
    "                    )\n",
    "                )\n",
    "            )\n",
    "        # remove the transformation due to the rest pose\n",
    "        G = G - self.pack(\n",
    "            np.matmul(\n",
    "                G,\n",
    "                np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])\n",
    "            )\n",
    "        )\n",
    "        # transformation of each vertex\n",
    "        T = np.tensordot(self.weights, G, axes=[[1], [0]])\n",
    "        rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))\n",
    "        v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3]\n",
    "        self.verts = v + self.trans.reshape([1, 3])\n",
    "        self.G = G\n",
    "\n",
    "    def rodrigues(self, r):\n",
    "        \"\"\"\n",
    "        Rodrigues' rotation formula that turns axis-angle vector into rotation\n",
    "        matrix in a batch-ed manner.\n",
    "    \n",
    "        Parameter:\n",
    "        ----------\n",
    "        r: Axis-angle rotation vector of shape [batch_size, 1, 3].\n",
    "    \n",
    "        Return:\n",
    "        -------\n",
    "        Rotation matrix of shape [batch_size, 3, 3].\n",
    "    \n",
    "        \"\"\"\n",
    "        theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)\n",
    "        # avoid zero divide\n",
    "        theta = np.maximum(theta, np.finfo(np.float64).tiny)\n",
    "        r_hat = r / theta\n",
    "        cos = np.cos(theta)\n",
    "        z_stick = np.zeros(theta.shape[0])\n",
    "        m = np.dstack([\n",
    "            z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1],\n",
    "            r_hat[:, 0, 2], z_stick, -r_hat[:, 0, 0],\n",
    "            -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick]\n",
    "        ).reshape([-1, 3, 3])\n",
    "        i_cube = np.broadcast_to(\n",
    "            np.expand_dims(np.eye(3), axis=0),\n",
    "            [theta.shape[0], 3, 3]\n",
    "        )\n",
    "        A = np.transpose(r_hat, axes=[0, 2, 1])\n",
    "        B = r_hat\n",
    "        dot = np.matmul(A, B)\n",
    "        R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m\n",
    "        return R\n",
    "\n",
    "    def with_zeros(self, x):\n",
    "        \"\"\"\n",
    "        Append a [0, 0, 0, 1] vector to a [3, 4] matrix.\n",
    "    \n",
    "        Parameter:\n",
    "        ---------\n",
    "        x: Matrix to be appended.\n",
    "    \n",
    "        Return:\n",
    "        ------\n",
    "        Matrix after appending of shape [4,4]\n",
    "    \n",
    "        \"\"\"\n",
    "        return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))\n",
    "\n",
    "    def pack(self, x):\n",
    "        \"\"\"\n",
    "        Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched\n",
    "        manner.\n",
    "    \n",
    "        Parameter:\n",
    "        ----------\n",
    "        x: Matrices to be appended of shape [batch_size, 4, 1]\n",
    "    \n",
    "        Return:\n",
    "        ------\n",
    "        Matrix of shape [batch_size, 4, 4] after appending.\n",
    "    \n",
    "        \"\"\"\n",
    "        return np.dstack((np.zeros((x.shape[0], 4, 3)), x))\n",
    "\n",
    "    def save_to_obj(self, path):\n",
    "        \"\"\"\n",
    "        Save the SMPL model into .obj file.\n",
    "    \n",
    "        Parameter:\n",
    "        ---------\n",
    "        path: Path to save.\n",
    "    \n",
    "        \"\"\"\n",
    "        with open(path, 'w') as fp:\n",
    "            for v in self.verts:\n",
    "                fp.write('v %f %f %f\\n' % (v[0], v[1], v[2]))\n",
    "            for f in self.faces + 1:\n",
    "                fp.write('f %d %d %d\\n' % (f[0], f[1], f[2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# load SMPL model and set to standard pose/shape\n",
    "smpl = SMPLModel('../SmplDownsampling/data/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl')\n",
    "trans = np.zeros(smpl.trans_shape)\n",
    "beta = np.zeros(smpl.beta_shape)\n",
    "pose = np.zeros(smpl.pose_shape)\n",
    "vert = smpl.set_params(beta=beta, pose=pose, trans=trans)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# calculate face areas\n",
    "smpl_face_area = []\n",
    "for f in smpl.faces:\n",
    "    v0, v1, v2 = vert[f[0]], vert[f[1]], vert[f[2]]\n",
    "    e0 = np.linalg.norm(v0 - v1)\n",
    "    e1 = np.linalg.norm(v1 - v2)\n",
    "    e2 = np.linalg.norm(v2 - v0)\n",
    "    s = np.sqrt((e0+e1+e2) * (e0+e1-e2) * (e1+e2-e0) * (e2+e0-e1)) * 0.25\n",
    "    smpl_face_area.append(s)\n",
    "np.save('./output/smpl_face_area.npy', smpl_face_area)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with open('./output/whole_body.obj', 'w') as fp:\n",
    "    for v in smpl.verts:\n",
    "        fp.write('v %f %f %f\\n' % (v[0], v[1], v[2]))\n",
    "\n",
    "# label some specific body parts according to skinning weights\n",
    "def get_vertices_bound_to_jnts(skinning_weights, jnts):\n",
    "    weights_of_interest = np.sum(skinning_weights[:, jnts], axis=1)\n",
    "    return np.where(weights_of_interest > 0.5)\n",
    "\n",
    "def save_part_of_smpl(smpl, vert_inds, filename):\n",
    "    with open(filename, 'w') as fp:\n",
    "        for vi in vert_inds:\n",
    "            v = smpl.verts[vi]\n",
    "            fp.write('v %f %f %f\\n' % (v[0], v[1], v[2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "left_hand_indices = get_vertices_bound_to_jnts(smpl.weights, [20, 22])[0]\n",
    "save_part_of_smpl(smpl, left_hand_indices, './output/left_hand.obj')\n",
    "np.save('./output/left_hand_indices.npy', left_hand_indices)\n",
    "right_hand_indices = get_vertices_bound_to_jnts(smpl.weights, [21, 23])[0]\n",
    "save_part_of_smpl(smpl, right_hand_indices, './output/right_hand.obj')\n",
    "np.save('./output/right_hand_indices.npy', right_hand_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "left_arm_indices = get_vertices_bound_to_jnts(smpl.weights, [16, 18])[0]\n",
    "save_part_of_smpl(smpl, left_arm_indices, './output/left_arm.obj')\n",
    "np.save('./output/left_arm_indices.npy', left_arm_indices)\n",
    "right_arm_indices = get_vertices_bound_to_jnts(smpl.weights, [17, 19])[0]\n",
    "save_part_of_smpl(smpl, right_arm_indices, './output/right_arm.obj')\n",
    "np.save('./output/right_arm_indices.npy', right_arm_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "head_indices = get_vertices_bound_to_jnts(smpl.weights, [12, 15])[0]\n",
    "save_part_of_smpl(smpl, head_indices, './output/head.obj')\n",
    "np.save('./output/head_indices.npy', head_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "left_leg_indices = get_vertices_bound_to_jnts(smpl.weights, [1, 4])[0]\n",
    "save_part_of_smpl(smpl, left_leg_indices, './output/left_leg.obj')\n",
    "np.save('./output/left_leg_indices.npy', left_leg_indices)\n",
    "right_leg_indices = get_vertices_bound_to_jnts(smpl.weights, [2, 5])[0]\n",
    "save_part_of_smpl(smpl, right_leg_indices, './output/right_leg.obj')\n",
    "np.save('./output/right_leg_indices.npy', right_leg_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "left_foot_indices = get_vertices_bound_to_jnts(smpl.weights, [7, 10])[0]\n",
    "save_part_of_smpl(smpl, left_foot_indices, './output/left_foot.obj')\n",
    "np.save('./output/left_foot_indices.npy', left_foot_indices)\n",
    "right_foot_indices = get_vertices_bound_to_jnts(smpl.weights, [8, 11])[0]\n",
    "save_part_of_smpl(smpl, right_foot_indices, './output/right_foot.obj')\n",
    "np.save('./output/right_foot_indices.npy', right_foot_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "torso_indices = get_vertices_bound_to_jnts(smpl.weights, [0, 3, 6, 9, 13, 14])[0]\n",
    "save_part_of_smpl(smpl, torso_indices, './output/torso.obj')\n",
    "np.save('./output/torso_indices.npy', torso_indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
